# -*- coding: utf-8 -*-

import psycopg2
import psycopg2.extras

from sqlexecutor import SqlExecutor


class PostgreSQL(SqlExecutor):

    def _init(self):
        super(PostgreSQL, self)._init()

        self.__connect()

    def close(self):
        try:
            self.__conn.close()
        except:
            pass

    def reset(self):
        self.close()

        self.__connect()
        self.sql = ''
        self.empty_params()

    def begin_transaction(self):
        if self.__status_in_transaction():
            raise Exception('A transaction is currently in progress.')

        self.__get_conn().autocommit = False

    def commit(self):
        if not self.__status_in_transaction():
            raise Exception('No transaction in progress.')

        self.__get_conn().commit()
        self.__get_conn().autocommit = True

    def rollback(self):
        if not self.__status_in_transaction():
            raise Exception('No transaction in progress.')

        self.__get_conn().rollback()
        self.__get_conn().autocommit = True

    def execute_non_query(self):
        self.__execute(fetch_result=False)

    def execute_scalar(self):
        ((ret,),) = self.__execute()

        return ret

    def execute_dataset(self):
        return self.__execute(dict_result=True)

    def __execute(self, fetch_result=True, dict_result=False):
        data = None

        with self.__get_conn().cursor(cursor_factory=psycopg2.extras.RealDictCursor if dict_result else None) as cur:
            cur.execute(self.sql, self.params)
            if fetch_result:
                data = cur.fetchall()

        return data

    def __status_in_transaction(self):
        return self.__get_conn().status == psycopg2.extensions.STATUS_IN_TRANSACTION

    def __get_conn(self):
        # if not self.__conn or self.__conn.closed != 0:
        #     self.__connect()

        return self.__conn

    def __connect(self):
        self.__conn = psycopg2.connect(**self._cfg['conn_args'])
        self.__conn.set_client_encoding('UTF8')
        self.__conn.autocommit = True
